package com.gitee.sidihuo.utils.async;

import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

@Slf4j
public class AsyncExecutor<P, R> {

    private Map<String, AsyncWorkerFunction<P, R>> workers = new HashMap<>();
    private Map<String, P> paras = new HashMap<>();

    public void addWorker(String workerKey, AsyncWorkerFunction<P, R> workerFunction, P para) {
        workers.put(workerKey, workerFunction);
        paras.put(workerKey, para);
    }

    public Map<String, R> executeAsync() {
        return executeAsync(1, System.currentTimeMillis() + "");
    }

    public Map<String, R> executeAsync(int maxTimeWaitMinutes, String traceId) {
        log.info("AsyncExecutor execute begin traceId = {} taskNum = {}", traceId, workers.size());
        CountDownLatch countDownLatch = new CountDownLatch(workers.size());
        Map<String, R> results = new HashMap<>();
        Iterator<Map.Entry<String, AsyncWorkerFunction<P, R>>> iterator = workers.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, AsyncWorkerFunction<P, R>> next = iterator.next();
            String workerKey = next.getKey();
            AsyncWorkerFunction<P, R> workerFunction = next.getValue();
            P para = paras.get(workerKey);
            AsyncWorkerThreadPool.getThreadPool().execute(new Runnable() {
                @Override
                public void run() {
                    long currentTimeMillisB = System.currentTimeMillis();
                    String subTrace = traceId + " " + workerKey;
                    log.info("AsyncExecutor execute begin sub thread subTrace = {}", subTrace);
                    R result = null;
                    try {
                        result = workerFunction.workAsync(para);
                    } catch (Exception e) {
                        log.warn("AsyncExecutor workerFunction workAsync failed", e);
                    }
                    results.put(workerKey, result);
                    countDownLatch.countDown();
                    long currentTimeMillisE = System.currentTimeMillis();
                    log.info("AsyncExecutor execute end sub thread subTrace = {} time = {}", subTrace, currentTimeMillisE - currentTimeMillisB);
                }
            });
        }

        long currentTimeMillis = System.currentTimeMillis();
        boolean await = false;
        try {
            await = countDownLatch.await(maxTimeWaitMinutes, TimeUnit.MINUTES);
        } catch (InterruptedException e) {
            log.warn("AsyncExecutor execute countDownLatch await InterruptedException", e);
        }
        long currentTimeMillis2 = System.currentTimeMillis();
        long totalTime = currentTimeMillis2 - currentTimeMillis;
        if (await) {
            log.info("AsyncExecutor execute countDownLatch finished traceId = {} totalTime = {}", traceId, totalTime);
        } else {
            log.warn("AsyncExecutor execute countDownLatch failed traceId = {} totalTime = {}", traceId, totalTime);
        }

        return results;
    }

}
